# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: jhkim

import json
import logging
from typing import List, Union
from collections import namedtuple

from tornado import gen
import zlib

import coin.flow.topic_map as topic_map
from coin.base.param_util import to_list
from coin.exchange.base.ws.subscriber import WebsocketJsonSubscriberBase

HashTimestamp = namedtuple('HashTimestamp', ['message_hash', 'timestamp'])

BOOK_UPDATE_INTERVAL_THRESHOLD = 20e9  # nanosecond


def _gen_message_hash(json_obj):
  data = json.dumps(json_obj)
  return zlib.crc32(data.encode())


def get_native_symbol(json_obj):
  if 'code' in json_obj:
    native_symbol = json_obj['code']
  elif 'cd' in json_obj:
    native_symbol = json_obj['cd']
  else:
    raise ValueError(json_obj)
  return native_symbol


def get_event_type(json_obj):
  if 'type' in json_obj:
    event_type = json_obj['type']
  elif 'ty' in json_obj:
    event_type = json_obj['ty']
  else:
    raise ValueError(json_obj)
  return event_type


def gen_subscribe_request_str(symbols):
  lst = []
  lst.append(dict(ticket="UNIQUE_TICKET",))
  lst.append(dict(
      type="orderbook",
      codes=symbols,
  ))
  lst.append(dict(
      type="trade",
      codes=symbols,
  ))
  lst_str = json.dumps(lst, separators=(',', ':'))
  return lst_str


class UpbitSubscriber(WebsocketJsonSubscriberBase):
  def __init__(self,
               queue,
               group_id: str,
               worker_id: str,
               symbol: Union[List[str], str],
               close_callback=None,
               logger=None):
    super().__init__(close_callback, logger or logging.getLogger(__name__))
    self._symbol = to_list(symbol)

    self._register_url('wss://api.upbit.com/websocket/v1')
    self._publisher = None
    self._subscribed = False
    self._init_publisher(queue, group_id, worker_id)
    self._request_str = gen_subscribe_request_str(self._symbol)
    self._last_hash_timestamp = {}

  @gen.coroutine
  def on_open(self):
    self._publish_topic_event('START')
    self.subscribe()

  def subscribe(self):
    assert not self._subscribed
    self.request()
    self._subscribed = True

  def request(self):
    self._ws.write_message(self._request_str)

  def _init_publisher(self, queue, group_id: str, worker_id: str):
    topic_depth = topic_map.search_topic_id('upbit_ws_depth_%s%s' % (group_id, worker_id))[0]
    topic_trade = topic_map.search_topic_id('upbit_ws_trade_%s%s' % (group_id, worker_id))[0]

    self._publisher = {
        'orderbook': queue.get_writer(topic_id=topic_depth[1]),
        'trade': queue.get_writer(topic_id=topic_trade[1]),
    }

    self._enable_topic_event_publisher(queue,
                                       'upbit_ws_topic_event', [topic_depth[0], topic_trade[0]])
    self._add_teardown_callback(self._teardown_publisher)

  def _teardown_publisher(self):
    if self._publisher:
      for _, publisher in self._publisher.items():
        publisher.close()
      self._publisher = None

  def on_message(self, timestamp: int, message: dict):
    if self.stopped:
      return
    if message is None:
      return
    event_type = get_event_type(message)
    if event_type == 'orderbook':
      if not self.is_book_update_valid(timestamp, message):
        logging.error('Book update interval longer than %d sec, restart...' %
                      (BOOK_UPDATE_INTERVAL_THRESHOLD / 1e9))
        self.stop(trigger_close_callback=True)
        return
    try:
      self._publisher[event_type].write(timestamp, message)
    except KeyError:
      self._logger.error('Unknown event_type: %s' % event_type)

  # In Upbit, when feed is normal, usually book will be updated in every
  # second. For small coin, it could be every 5 - 10 seconds.
  # When feed is abnormal, book update is typically every 25 seconds. And
  # book value will be the same for that period and usually the value is
  # incorrect.
  # Here check the book content and update timestamp, if content hash is the
  # the same and it's updated after more than 20 seconds, consider shit happens
  # and restart the subscriber.
  def is_book_update_valid(self, timestamp, json_obj):
    native_symbol = get_native_symbol(json_obj)
    last_hash_timestamp = self._last_hash_timestamp.get(native_symbol, None)

    if last_hash_timestamp is not None:
      message_hash = _gen_message_hash(json_obj)
      valid = ((message_hash != last_hash_timestamp.message_hash)
               or (timestamp - last_hash_timestamp.timestamp <= BOOK_UPDATE_INTERVAL_THRESHOLD))
      self._last_hash_timestamp[native_symbol] = HashTimestamp(message_hash, timestamp)
      return valid
    else:
      self._last_hash_timestamp[native_symbol] = HashTimestamp(_gen_message_hash(json_obj),
                                                               timestamp)
      return True
